package graph

type NodeRecord struct {
	Node *Node
	Distance int
}

func DijKstraByHeap(from *Node, size int) map[*Node]int {
	distanceMap := map[*Node]int{}
	mySmallHeap := NewMySmallHeap(size)
	mySmallHeap.InsertOrUpdateOrIgnore(from, 0)
	for !mySmallHeap.IsEmpty() {
		nodeRecord := mySmallHeap.Pop()
		cur := nodeRecord.Node
		distance := nodeRecord.Distance
		for _, edge := range cur.OutEdges {
			mySmallHeap.InsertOrUpdateOrIgnore(edge.To, edge.Weight + distance)
		}
		distanceMap[cur] =  distance
	}
	return distanceMap
}

func NewMySmallHeap (size int) MySmallHeap {
	return MySmallHeap{
		nodes: make([]*Node, size),
		NodeIndex: make(map[*Node]int),
		DistanceMap: make(map[*Node]int),
	}
}


type MySmallHeap struct {
	nodes []*Node
	size  int
	NodeIndex map[*Node]int
	DistanceMap map[*Node]int
}

func (m *MySmallHeap) IsEmpty() bool {
	return m.size == 0
}

func (m *MySmallHeap) InsertOrUpdateOrIgnore(node *Node, distance int) {
	if m.InHeap(node) {
		m.DistanceMap[node] = min(m.DistanceMap[node], distance)
		m.InsertHeapify(m.NodeIndex[node])
	}
	if !m.IsEntered(node) {
		m.DistanceMap[node] = distance
		m.nodes[m.size] = node
		m.NodeIndex[node] = m.size
		m.size++
		m.InsertHeapify(m.size-1)

	}
}

func (m *MySmallHeap) IsEntered (node *Node) bool {
	return m.Contains(node)
}

func  (m *MySmallHeap) Pop() NodeRecord {
	ans := NodeRecord{m.nodes[0], m.DistanceMap[m.nodes[0]]}
	m.Swap(0, m.size-1)
	m.NodeIndex[m.nodes[m.size-1]] = -1
	delete(m.DistanceMap, m.nodes[m.size-1])
	m.size--
	m.Heapify(0)
	return ans
}

func (m *MySmallHeap) InsertHeapify(index int) {
	for m.DistanceMap[m.nodes[index]] < m.DistanceMap[m.nodes[(index-1)/2]] {
		m.Swap(index, (index-1)/2)
		index = (index-1)/2
	}
}

func (m *MySmallHeap) Heapify(index int) {
	left := 2*index + 1
	for left < m.size {
		small := left
		if left+1 < m.size && m.DistanceMap[m.nodes[left]] > m.DistanceMap[m.nodes[left+1]] {
			small = left + 1
		}
		if m.DistanceMap[m.nodes[small]] > m.DistanceMap[m.nodes[index]] {
			break
		}
		m.Swap(index, small)
		index = small
		left = 2*index + 1
	}
}

func (m *MySmallHeap) Insert (index int) {
	parent := (index-1)/2
	for parent >= 0 {
		if m.DistanceMap[m.nodes[index]] <  m.DistanceMap[m.nodes[parent]] {
			m.Swap(index, parent)
		} else {
			break
		}
		index = parent
		parent = (index-1)/2
	}
}

func (m *MySmallHeap) Swap(i, j int) {
	m.NodeIndex[m.nodes[i]], m.NodeIndex[m.nodes[j]] = j, i
	m.nodes[i], m.nodes[j] = m.nodes[j], m.nodes[i]
}

func (m *MySmallHeap) Contains(node *Node) bool {
	if _, ok :=m.NodeIndex[node]; ok {
		return true
	}
	return false
}

func (m *MySmallHeap) InHeap(node *Node) bool {
	return m.Contains(node) && m.NodeIndex[node] != -1
}